"""
批量数据操作功能演示示例

该示例展示了SQLServerManager类的批量数据处理功能，包括批量插入、批量更新和批量Upsert操作，
特别适合处理大量数据时提高效率。
"""
from modules.database.sqlserver_manager import SQLServerManager
import logging
import random
from faker import Faker

# 配置日志
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger = logging.getLogger('sqlserver_manager')
logger.addHandler(handler)
logger.setLevel(logging.INFO)

# 初始化Faker生成测试数据
fake = Faker('zh_CN')


def generate_test_data(count: int) -> List[Dict[str, Any]]:
    """生成测试数据"""
    departments = ['Engineering', 'Marketing', 'HR', 'Finance', 'Sales']
    data = []
    for _ in range(count):
        data.append({
            'Name': fake.name(),
            'Department': random.choice(departments),
            'Salary': round(random.uniform(50000, 150000), 2),
            'HireDate': fake.date_between(start_date='-5y', end_date='today')
        })
    return data

def demonstrate_bulk_operations():
    # 数据库连接参数
    db_config = {
        'server': 'localhost',
        'database': 'test_db',
        'username': 'sa',
        'password': 'your_password',
        'port': 1433
    }

    # 使用上下文管理器确保连接正确关闭
    with SQLServerManager(**db_config) as db:
        try:
            # 1. 创建演示表
            logger.info("创建演示表...")
            db.execute_non_query("""
                CREATE TABLE IF NOT EXISTS BulkTestEmployees (
                    ID INT PRIMARY KEY IDENTITY(1,1),
                    Name NVARCHAR(50) NOT NULL,
                    Department NVARCHAR(50) NOT NULL,
                    Salary DECIMAL(10,2) NOT NULL,
                    HireDate DATE NOT NULL
                )
            """)

            # 2. 批量插入示例
            logger.info("\n=== 演示批量插入 ===")
            # 生成1000条测试数据
            test_data = generate_test_data(1000)
            logger.info(f"生成 {len(test_data)} 条测试数据")
            
            # 执行批量插入
            affected_rows = db.bulk_insert("BulkTestEmployees", test_data, batch_size=200)
            logger.info(f"批量插入完成，影响 {affected_rows} 行")

            # 3. 批量更新示例
            logger.info("\n=== 演示批量更新 ===")
            # 查询需要更新的数据
            employees = db.execute_query("SELECT TOP 200 ID, Salary FROM BulkTestEmployees WHERE Department = 'Engineering'")
            logger.info(f"找到 {len(employees)} 名工程部员工需要加薪")
            
            # 准备批量更新数据
            update_data = []
            for emp in employees:
                # 加薪10%
                new_salary = round(emp['Salary'] * 1.10, 2)
                update_data.append((
                    {'Salary': new_salary},
                    'ID = ?',
                    (emp['ID'],)
                ))
            
            # 执行批量更新
            affected_rows = db.bulk_update("BulkTestEmployees", update_data, batch_size=100)
            logger.info(f"批量更新完成，影响 {affected_rows} 行")

            # 4. 批量Upsert示例 (更新现有记录，插入新记录)
            logger.info("\n=== 演示批量Upsert ===")
            # 创建混合了现有和新记录的数据
            existing_employees = db.execute_query("SELECT TOP 50 ID, Name, Department, Salary, HireDate FROM BulkTestEmployees")
            new_employees = generate_test_data(50)
            
            # 修改部分现有员工数据
            upsert_data = []
            for emp in existing_employees:
                # 随机调整薪资
                emp['Salary'] = round(emp['Salary'] * random.uniform(0.95, 1.05), 2)
                upsert_data.append(emp)
            
            # 添加新员工数据
            upsert_data.extend(new_employees)
            logger.info(f"准备 {len(upsert_data)} 条Upsert数据（{len(existing_employees)}条更新，{len(new_employees)}条新增）")
            
            # 执行批量Upsert
            affected_rows = db.bulk_upsert(
                "BulkTestEmployees",
                upsert_data,
                merge_condition="target.ID = source.ID",
                batch_size=100
            )
            logger.info(f"批量Upsert完成，影响 {affected_rows} 行")

        except Exception as e:
            logger.error(f"操作发生错误: {str(e)}", exc_info=True)
        finally:
            # 5. 清理演示表
            logger.info("\n清理演示表...")
            db.execute_non_query("DROP TABLE IF EXISTS BulkTestEmployees")

if __name__ == "__main__":
    demonstrate_bulk_operations()